#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import numpy as np
from osgeo import gdal, osr
import trimesh

def convert_osgb_to_ortho(input_osgb, output_dom, output_dsm, resolution=0.01):
    """
    将OSGB文件转换为正射影像（DOM）和数字表面模型（DSM）
    
    参数:
        input_osgb (str): 输入OSGB文件的路径
        output_dom (str): 输出DOM文件的路径
        output_dsm (str): 输出DSM文件的路径
        resolution (float): 输出影像的分辨率（米/像素），默认0.01
    """
    try:
        # 1. 读取OSGB文件并转换为点云
        print(f"正在读取OSGB文件: {input_osgb}")
        mesh = trimesh.load(input_osgb)
        
        # 获取顶点和颜色
        points = mesh.vertices
        if hasattr(mesh.visual, 'vertex_colors'):
            colors = mesh.visual.vertex_colors[:, :3] / 255.0  # 转换为0-1范围
        else:
            # 如果没有颜色信息，使用默认颜色
            colors = np.ones_like(points) * 0.5
        
        # 3. 计算边界框
        min_x, max_x = np.min(points[:, 0]), np.max(points[:, 0])
        min_y, max_y = np.min(points[:, 1]), np.max(points[:, 1])
        min_z, max_z = np.min(points[:, 2]), np.max(points[:, 2])
        
        # 4. 计算输出影像的大小
        width = int((max_x - min_x) / resolution)
        height = int((max_y - min_y) / resolution)
        
        print(f"输出影像大小: {width}x{height} 像素")
        
        # 5. 创建DOM和DSM数组
        dom = np.zeros((height, width, 3), dtype=np.uint8)
        dsm = np.full((height, width), -9999, dtype=np.float32)  # 使用-9999表示无数据
        
        # 6. 投影点云到正射平面
        for i in range(len(points)):
            x, y, z = points[i]
            color = colors[i] * 255  # 转换为0-255范围
            
            # 计算像素坐标
            pixel_x = int((x - min_x) / resolution)
            pixel_y = int((max_y - y) / resolution)  # 注意y轴方向
            
            if 0 <= pixel_x < width and 0 <= pixel_y < height:
                # 更新DOM（只在当前点比已有点更高时更新）
                if z > dsm[pixel_y, pixel_x] or dsm[pixel_y, pixel_x] == -9999:
                    dom[pixel_y, pixel_x] = color
                    dsm[pixel_y, pixel_x] = z
        
        # 7. 填充DOM中的空洞（简单的邻近像素平均）
        kernel_size = 3
        for y in range(height):
            for x in range(width):
                if not dom[y, x].any():  # 如果是黑色像素
                    y1 = max(0, y - kernel_size//2)
                    y2 = min(height, y + kernel_size//2 + 1)
                    x1 = max(0, x - kernel_size//2)
                    x2 = min(width, x + kernel_size//2 + 1)
                    
                    patch = dom[y1:y2, x1:x2]
                    valid_pixels = patch[patch.any(axis=2)]
                    if len(valid_pixels) > 0:
                        dom[y, x] = np.mean(valid_pixels, axis=0)
        
        # 8. 保存DOM
        print(f"正在保存DOM: {output_dom}")
        driver = gdal.GetDriverByName('GTiff')
        dataset = driver.Create(output_dom, width, height, 3, gdal.GDT_Byte)
        
        # 设置地理参考
        geotransform = [min_x, resolution, 0, max_y, 0, -resolution]
        dataset.SetGeoTransform(geotransform)
        
        # 设置投影
        srs = osr.SpatialReference()
        srs.ImportFromEPSG(4326)  # WGS84
        dataset.SetProjection(srs.ExportToWkt())
        
        # 写入数据
        for i in range(3):
            dataset.GetRasterBand(i + 1).WriteArray(dom[:,:,i])
        dataset = None
        
        # 9. 保存DSM
        print(f"正在保存DSM: {output_dsm}")
        dataset = driver.Create(output_dsm, width, height, 1, gdal.GDT_Float32)
        dataset.SetGeoTransform(geotransform)
        dataset.SetProjection(srs.ExportToWkt())
        dataset.GetRasterBand(1).WriteArray(dsm)
        dataset.GetRasterBand(1).SetNoDataValue(-9999)
        dataset = None
        
        print("处理完成！")
        return True
        
    except Exception as e:
        print(f"处理过程中出现错误: {str(e)}")
        return False

def main():
    if len(sys.argv) != 4:
        print("使用方法: python osgb_to_ortho.py <输入OSGB文件> <输出DOM文件> <输出DSM文件>")
        print("示例: python osgb_to_ortho.py input.osgb output_dom.tif output_dsm.tif")
        sys.exit(1)
    
    input_osgb = sys.argv[1]
    output_dom = sys.argv[2]
    output_dsm = sys.argv[3]
    
    if not os.path.exists(input_osgb):
        print(f"错误：输入文件 {input_osgb} 不存在")
        sys.exit(1)
    
    success = convert_osgb_to_ortho(input_osgb, output_dom, output_dsm)
    if not success:
        sys.exit(1)

if __name__ == "__main__":
    main() 